package ratelimiter

import (
	"context"
	"errors"

	"github.com/go-kratos/kratos/v2/middleware"
	"golang.org/x/time/rate"
)

func NewRatelimitMiddleware(rateLimit, burst int) middleware.Middleware {
	limiter := rate.NewLimiter(rate.Limit(rateLimit), burst)
	return func(handler middleware.Handler) middleware.Handler {
		return func(ctx context.Context, req interface{}) (interface{}, error) {
			if !limiter.Allow() {
				return nil, errors.New("too many requests you have made")
			}
			return handler(ctx, req)
		}
	}
}
